#include <iostream>
#include <array>
#include <ctime>
#include <cstring>

#ifdef WAFER_OPENSSL
#include <openssl/evp.h>
#endif

#include "container.h"

#define HASH_SIZE 20
#define MSG_LEN 2
#define MSG_NUM 100
#define TOTAL_MSG_SIZE MSG_LEN * MSG_NUM

#ifndef WAFER_OPENSSL
extern "C" {
    void CUDA(
    MemRef<unsigned char, 1> *hash, MemRef<unsigned char, 1> *input, 
    int num, int len, MemRef<unsigned char, 1> *res);
}

void SHA1_CRACK_OP(MemRef<unsigned char, 1> *hash, MemRef<unsigned char, 1> *input, 
                   int num, int len, MemRef<unsigned char, 1> *res) {
    CUDA(hash, input, num, len, res);
}
#else
#define OPENSSL_CPU_NUM 2
extern "C" uint64_t sstGetCpuId();
#endif

// Get cpu id in sst.
// extern "C" uint64_t sstGetCpuId();

uint8_t candidates[TOTAL_MSG_SIZE] = {'0', '0', '0', '1', '0', '2', '0', '3', '0', '4', '0', '5', '0', '6', '0', '7', '0', '8', '0', '9',
                                      '1', '0', '1', '1', '1', '2', '1', '3', '1', '4', '1', '5', '1', '6', '1', '7', '1', '8', '1', '9',
                                      '2', '0', '2', '1', '2', '2', '2', '3', '2', '4', '2', '5', '2', '6', '2', '7', '2', '8', '2', '9',
                                      '3', '0', '3', '1', '3', '2', '3', '3', '3', '4', '3', '5', '3', '6', '3', '7', '3', '8', '3', '9',
                                      '4', '0', '4', '1', '4', '2', '4', '3', '4', '4', '4', '5', '4', '6', '4', '7', '4', '8', '4', '9',
                                      '5', '0', '5', '1', '5', '2', '5', '3', '5', '4', '5', '5', '5', '6', '5', '7', '5', '8', '5', '9',
                                      '6', '0', '6', '1', '6', '2', '6', '3', '6', '4', '6', '5', '6', '6', '6', '7', '6', '8', '6', '9',
                                      '7', '0', '7', '1', '7', '2', '7', '3', '7', '4', '7', '5', '7', '6', '7', '7', '7', '8', '7', '9',
                                      '8', '0', '8', '1', '8', '2', '8', '3', '8', '4', '8', '5', '8', '6', '8', '7', '8', '8', '8', '9',
                                      '9', '0', '9', '1', '9', '2', '9', '3', '9', '4', '9', '5', '9', '6', '9', '7', '9', '8', '9', '9'};
uint8_t target[HASH_SIZE] = {0xeb, 0x4a, 0xc3, 0x03, 0x3e, 0x8a, 0xb3, 0x59, 0x1e, 0x0f, 0xce, 0xfa, 0x8c, 0x26, 0xce, 0x3f, 0xd3, 0x6d, 0x5a, 0x0f};
uint8_t result[MSG_LEN] = {0};
MemRef<unsigned char, 1> input{candidates, TOTAL_MSG_SIZE};
MemRef<unsigned char, 1> hash{target, HASH_SIZE};
MemRef<unsigned char, 1> output{result, MSG_LEN};


void printMatches(const char* title, uint8_t* match) {
  // Print the output.
  std::cout << std::endl;
  std::cout << title;
  std::cout << "[ ";
  for (size_t i = 0; i < MSG_LEN; i++) {
    printf("%c ", match[i]);
  }
  std::cout << "]";
  std::cout << std::endl;
}

void mlirCrackSHA1() {
#ifdef WAFER_OPENSSL
  uint64_t cpuid = sstGetCpuId();
  uint64_t step = MSG_NUM / OPENSSL_CPU_NUM;
  uint64_t start = cpuid * step;
  uint64_t end = (cpuid == OPENSSL_CPU_NUM - 1) ? MSG_NUM : ((cpuid + 1) * step);

  const clock_t begin_time = clock();
  uint8_t *tmp = nullptr;
  const EVP_MD* sha1 = EVP_sha1();
  for (int i = start; i < end; i++) {
    tmp = static_cast<uint8_t *>(malloc(HASH_SIZE * sizeof(uint8_t)));

    uint8_t* msgBase = candidates + (i * MSG_LEN);
    EVP_MD_CTX* sha1_ctx = EVP_MD_CTX_new();
    EVP_DigestInit_ex(sha1_ctx, sha1, nullptr);
    EVP_DigestUpdate(sha1_ctx, msgBase, MSG_LEN);
    uint32_t size = HASH_SIZE;
    EVP_DigestFinal_ex(sha1_ctx, tmp, &size);

    if (memcmp(target, tmp, MSG_LEN) == 0) {
      printMatches("OpenSSL SHA1 Crack: ", msgBase);
    }

    free(tmp);
  }

  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  std::cout << "cpuid: " << cpuid << std::endl << "handle data: [" << start << ", " << end << "]" << std::endl <<"Calculate duration: " << std::fixed << duration << std::endl;
#else
  const clock_t begin_time = clock();
  #pragma wafer parallel
  SHA1_CRACK_OP(&hash, &input, MSG_NUM, MSG_LEN, &output);
  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  printMatches("MLIR SHA1 Crack: ", result);
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
#endif
}

int main(){
    mlirCrackSHA1();
    return 0;
}
